import akshare as ak
import pandas as pd
import finplot as fplt
import numpy as np


def get_stock_data(symbol="sh600036", period="daily", start_date="20250101", end_date="20250731", adjust="qfq"):
    """
    获取股票数据

    Parameters:
    symbol (str): 股票代码
    period (str): 时间周期 ("daily", "weekly", "monthly")
    start_date (str): 开始日期 (YYYYMMDD)
    end_date (str): 结束日期 (YYYYMMDD)
    adjust (str): 复权方式 ("qfq": 前复权, "hfq": 后复权, "": 不复权)

    Returns:
    pd.DataFrame: 股票数据
    """
    df = ak.stock_zh_a_hist(
        symbol=symbol,
        period=period,
        start_date=start_date,
        end_date=end_date,
        adjust=adjust
    )
    return df


def calculate_ene(df, N=10, M1=11, M2=9):
    """
    计算ENE轨道线指标
    ENE轨道线是基于移动平均线和一定比例的上下轨道线构成的技术分析指标
    
    Parameters:
    df (pd.DataFrame): 股票数据
    N (int): 移动平均周期
    M1 (int): 上轨道百分比
    M2 (int): 下轨道百分比
    
    Returns:
    pd.DataFrame: 添加了ENE指标的数据
    """
    df = df.copy()
    
    # 计算移动平均线
    df['MA'] = df['close'].rolling(window=N).mean()
    
    # 计算上轨道线 (UPPER)
    df['ENE_UPPER'] = (1 + M1 / 100) * df['MA']
    
    # 计算下轨道线 (LOWER)
    df['ENE_LOWER'] = (1 - M2 / 100) * df['MA']
    
    # 计算ENE中线 (ENE线)
    df['ENE'] = (df['ENE_UPPER'] + df['ENE_LOWER']) / 2
    
    return df

def calculate_moving_averages(df, windows=[5, 10, 20]):
    """
    计算价格均线
    
    Parameters:
    df (pd.DataFrame): 股票数据
    windows (list): 均线周期列表
    
    Returns:
    pd.DataFrame: 添加了均线的数据
    """
    df = df.copy()
    
    # 计算价格均线
    for window in windows:
        df[f'ma_price_{window}'] = df['close'].rolling(window=window).mean()
        
    return df

def calculate_pressure_lines(df, period=20):
    """
    计算近期压力线（基于近期最高价和最低价）
    
    Parameters:
    df (pd.DataFrame): 股票数据
    period (int): 计算周期
    
    Returns:
    tuple: (近期最高值, 近期最低值)
    """
    df = df.copy()
    
    # 计算近期最高价和最低价
    recent_high = df['high'].tail(period).max()
    recent_low = df['low'].tail(period).min()
    
    return recent_high, recent_low

def plot_stock_chart(df, symbol="600036"):
    """
    使用finplot绘制股票图表，包含K线、均线、ENE轨道线和压力线
    
    Parameters:
    df (pd.DataFrame): 股票数据
    symbol (str): 股票代码
    """
    # 只取最近100天的数据
    df = df.tail(100)
    
    # 计算压力线
    pressure_high, pressure_low = calculate_pressure_lines(df)
    
    # 创建图表
    ax1, ax2, ax3 = fplt.create_plot(f'{symbol} 股票分析图表', rows=3, init_zoom_periods=100)
    
    # 绘制K线图
    fplt.candlestick_ochl(df[['open', 'close', 'high', 'low']], ax=ax1)
    
    # 绘制价格均线
    fplt.plot(df['ma_price_5'], ax=ax1, color='#0000ff', legend='MA5')
    fplt.plot(df['ma_price_10'], ax=ax1, color='#ffa500', legend='MA10')
    fplt.plot(df['ma_price_20'], ax=ax1, color='#800080', legend='MA20')
    
    # 绘制ENE轨道线
    fplt.plot(df['ENE_UPPER'], ax=ax1, color='#ff0000', legend='ENE_UPPER')
    fplt.plot(df['ENE_LOWER'], ax=ax1, color='#008000', legend='ENE_LOWER')
    fplt.plot(df['ENE'], ax=ax1, color='#000000', legend='ENE')
    
    # 绘制压力线
    fplt.plot(pd.Series([pressure_high] * len(df)), ax=ax1, color='#ff0000', style='--', legend=f'压力线: {pressure_high:.2f}')
    fplt.plot(pd.Series([pressure_low] * len(df)), ax=ax1, color='#008000', style='--', legend=f'支撑线: {pressure_low:.2f}')
    
    # 绘制成交量
    fplt.volume_ocv(df[['open', 'close', 'volume']], ax=ax2)
    
    # 计算并绘制量均线
    df['ma_volume_5'] = df['volume'].rolling(window=5).mean()
    df['ma_volume_10'] = df['volume'].rolling(window=10).mean()
    fplt.plot(df['ma_volume_5'], ax=ax2, color='#ffa500', legend='Volume MA5')
    fplt.plot(df['ma_volume_10'], ax=ax2, color='#800080', legend='Volume MA10')
    
    # 简化的筹码分布（以收盘价的标准差表示）
    df['price_std'] = df['close'].rolling(window=20).std()
    df['chip_distribution'] = 1 / (1 + df['price_std'] / df['close'])
    
    # 绘制筹码分布
    fplt.plot(df['chip_distribution'], ax=ax3, color='#8B4513', legend='筹码集中度')
    fplt.set_y_range(0, 1.2, ax=ax3)
    
    # 显示图表
    fplt.show()

def main():
    # 获取股票数据
    print("正在获取股票数据...")
    df = get_stock_data()
    
    # 计算技术指标
    print("正在计算技术指标...")
    df = calculate_moving_averages(df)
    df = calculate_ene(df)
    
    # 显示数据信息
    print("\n数据基本信息:")
    print(df[['date', 'open', 'high', 'low', 'close', 'volume']].tail())
    
    # 绘制图表
    print("正在绘制图表...")
    plot_stock_chart(df)

if __name__ == "__main__":
    main()
